package com.tjc.hashing;

import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.*;

/**
 * @Author: tjc
 * @description
 * @Date Created in 15:35 2022-09-03
 */
public class ConsistentHashing {

    Map<String, ServerNode> servers = new HashMap<>();
    SortedMap<Integer, VNode> vNodes = new TreeMap<>((key1, key2) -> compareUnsigned(key1, key2));
    int vNodeCount;

    public ConsistentHashing(int vNodeCount) {
        this.vNodeCount = vNodeCount;
    }

    /**
     * copy from {@link java.util.HashMap#hash(Object)}
     */
    static final int hash(Object key) {
        // 用md5 更加均衡
        int h;
        return (key == null) ? 0 : (h = key.hashCode()) ^ (h >>> 16);
    }

    static final int compareUnsigned(int a, int b) {
        int fa = a >>> 31;
        int fb = b >>> 31;
        if ((fa | fb) == 0) return Integer.compare(a, b);
        if ((fa & fb) > 0) return Integer.compare(a, b);
        return fa > fb ? 1 : -1;
    }

    public void addServer(String server) {
        ServerNode e = new ServerNode(server);
        servers.put(server, e);

        for (int i = 0; i < vNodeCount; i++) {
            int h = hash();
            VNode n;
            if ((n = vNodes.get(h)) != null) {
                n.serverNodes.add(e);
            } else {
                n = new VNode(h, e);
            }
            e.vNodes.add(n);
            vNodes.put(n.hashKey, n);
        }
    }

    private static int hash() {
        MessageDigest instance = null;

        try {
            instance = MessageDigest.getInstance("MD5");
        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
        }
        byte[] digest = instance.digest();

        int h = 0;
        for (int i = 0; i < 4; i++) {
            h <<= 8;
            h |= ((int) digest[i]) & 0xFF;
        }
        return h;
    }

    public void removeServer(String server) {
        // todo
        ServerNode node = servers.remove(server);
        if (node != null) {
            node.vNodes.stream().forEach(e -> vNodes.remove(e.hashKey));
        }
    }

    /**
     * @return server
     */
    public Object access(String sourceId) {
        if (servers.size() == 0) return null;
        int h = hash(sourceId);

        SortedMap<Integer, VNode> map = vNodes.tailMap(h);
        VNode match;
        if (!map.isEmpty()) {
            match = map.get(map.firstKey());
        } else {
            match = vNodes.get(vNodes.firstKey());
        }
        int i = h % match.serverNodes.size();
        return match.serverNodes.get(i);
    }


    class ServerNode {
        int position;
        String server;
        SortedSet<VNode> vNodes = new TreeSet();

        public ServerNode(String server) {
            this.server = server;
            this.position = hash(server);
        }

        @Override
        public String toString() {
            return "ServerNode{" +
                    "position=" + Integer.toBinaryString(position) +
                    ", server='" + server + '\'' +
                    '}';
        }
    }

    class VNode {
        int hashKey;
        List<ServerNode> serverNodes = new ArrayList<>();

        public VNode(int h, ServerNode e) {
            this.hashKey = h;
            serverNodes.add(e);
        }
    }
}
